I think we tend to pin/unpin frequently enough that we can add build args to the env section in.github/workflows/build_docker_images.yml and build the docker images based on these. Changes needed

In the build images workflow

 env:
   REGISTRY: diffusers
   CI_SLACK_CHANNEL: ${{ secrets.CI_DOCKER_CHANNEL }}
   TORCH_VERSION: 2.10.0 # set to "" to use latest 

Update the build instructions

        with:
          no-cache: true
          context: ./docker/${{ matrix.image-name }}
          push: true
          tags: ${{ env.REGISTRY }}/${{ matrix.image-name }}:latest
          build-args: |
            TORCH_VERSION=${{ env.TORCH_VERSION }}

torchvision and torchaudio will resolve to the appropriate versions based on the torch version.

Inside the Dockerfile, we would need

RUN uv pip install --no-cache-dir \
    "torch${TORCH_VERSION:+==${TORCH_VERSION}}" \
    torchvision \
    torchaudio \
    --extra-index-url https://download.pytorch.org/whl/cpu

torchvision and torchaudio will resolve to the appropriate versions based on the torch version.

But when we do pip install torch (the current stage):

RUN uv pip install --no-cache-dir \
    torch \
    torchvision \
    torchaudio 

somehow the resolution isn't appropriate.

--extra-index-url https://download.pytorch.org/whl/cpu

This feels inappropriate for non-CPU envs i.e., for the GPU docker containers.